/**
 *@file main.cpp
 *@author lynsdk team
 *@version v1.0
 *@date 2022-09-06
 *@par Copyright:
 *© 2018 北京灵汐科技有限公司 版权所有。\n
 * 注意：以下内容均为北京灵汐科技有限公司原创，未经本公司允许，不得转载，否则将视为侵权；对于不遵守此声明或者其他违法使用以下内容者，本公司依法保留追究权。\n
 *© 2018 Lynxi Technologies Co., Ltd. All rights reserved. \n
 * NOTICE: All information contained here is, and remains the property of Lynxi.
 *This file can not be copied or distributed without the permission of Lynxi
 *Technologies Co., Ltd.
 *@brief This file describes the sample for using apu inference on lynxi KA200
 *apu
 */

#include <lyn_api.h>
#include <sys/stat.h>
#include <unistd.h>
#include <fstream>
#include <iomanip>
#include <iostream>
#include <memory>
#include <string>
#include "bufferPool.hpp"
#include "logging.h"
#include "util.hpp"

#include <semaphore.h>
#include <stdio.h>
#include <atomic>
#include <cassert>
#include <condition_variable>
#include <fstream>
#include <iomanip>
#include <iostream>
#include <mutex>
#include <queue>
#include <string>
#include <thread>
#include "CLI/CLI11.hpp" // 引入 CLI11 头文件

#define SKIP_FRAME_NUM (1)
#define SWAP_POOL_COUNT 2
#define FRAME_COUNT 1000
#define ERROR_RET(function)                      \
  do                                             \
  {                                              \
    if (0 != (err = (function)))                 \
    {                                            \
      std::cout << "error " << err << std::endl; \
      goto ret;                                  \
    }                                            \
  } while (0)

static std::string g_argModuleFile;
static bool g_streamError = false;
static int64_t g_frameCount = FRAME_COUNT;
static int g_deviceId = 0;
static int g_count = 0;
static std::string g_file_input;
static std::string g_file_output;
static int g_dynamicHeight = 0;
static int g_dynamicWidth = 0;
static int g_batch = 0;
// static uint32_t heightNow = 0;
// static uint32_t widthNow = 0;

typedef struct
{
  BufferPool *InPool;
  void *InBuffer;
  BufferPool *OutPool;
  void *OutBuffer;
} apu_cb_info;

lynError_t apu_buffer_callback(void *userData)
{
  lynError_t ret = 0;
  apu_cb_info *buffer = (apu_cb_info *)userData;

  if (!g_file_output.empty())
  {
    std::ofstream outputFile(g_file_output, std::ios::out | std::ios::binary);
    if (false == outputFile.is_open())
    {
      std::cout << "test open failed " << g_file_output.c_str();
    }
    else
    {
      uint64_t apuOutSize = buffer->OutPool->ElementSize();
      void *hostout = malloc(apuOutSize);
      if (nullptr == hostout)
      {
        std::cerr << "malloc failed " << apuOutSize;
      }
      else
      {
        lynMemcpy(hostout, buffer->OutBuffer, apuOutSize, ServerToClient);
        outputFile.write((char *)hostout, apuOutSize);
        outputFile.close();
        free(hostout);
      }
    }
  }

  buffer->InPool->Push(buffer->InBuffer);
  buffer->OutPool->Push(buffer->OutBuffer);
  free(buffer);
  g_count++;
  // printf("[ %d ] \r",g_count);
  return ret;
}

std::string getCurTimeStr()
{
  char buf[128] = {0};

  time_t tt = time(nullptr);
  struct tm *ptr = localtime(&tt);
  char strTime[64] = {0};
  strftime(strTime, sizeof(strTime), "%Y-%m-%d %H:%M:%S", ptr);

  struct timeval tmv;
  gettimeofday(&tmv, nullptr);
  sprintf(buf, "%s:%06ld", strTime, tmv.tv_usec);

  return std::string(buf);
}

double calcTFLOPS(lynModelDesc_t *modelDesc, double fps)
{
  uint64_t bignum = 30 * 2 * modelDesc->inputTensorAttrArray[0].dims[1] *
                    modelDesc->inputTensorAttrArray[0].dims[2] *
                    modelDesc->outputTensorAttrArray[0].dims[2];
  double tflops = fps * bignum / 1000000000000;

  return tflops;
}

int apuinfer()
{
  lynError_t err = 0;
  lynContext_t context = nullptr;
  lynStream_t streamSend = nullptr;
  lynStream_t streamRecv = nullptr;
  lynModelDesc_t *modelDesc = nullptr;
  lynModel_t model = 0;
  lynEvent_t eventStart = nullptr;
  lynEvent_t eventStop = nullptr;
  int frameNum = 0;
  double fps = 0;
  // double tflops = 0 ;
  float timees = 0.0;
  uint32_t batch = 0;
  BufferPool *apuIn = nullptr;
  BufferPool *apuOut = nullptr;
  int ApuOutSize, ApuInSize;
  void *apuInBuffer;
  void *apuOutBuffer;
  void *hostin = nullptr;

  std::chrono::steady_clock::time_point tFrameBeg, tFrameEnd;
  std::chrono::duration<double, std::milli> tFrameSpan;
  std::string start_time, end_time;

  ERROR_RET(lynCreateContext(&context, g_deviceId));
  ERROR_RET(lynRegisterErrorHandler(StreamErrorHandler, nullptr));

  ERROR_RET(lynCreateStream(&streamSend));
  ERROR_RET(lynCreateStream(&streamRecv));
  ERROR_RET(lynCreateEvent(&eventStart));
  ERROR_RET(lynCreateEvent(&eventStop));
  ERROR_RET(lynLoadModel(g_argModuleFile.c_str(), &model));
  ERROR_RET(lynModelGetDesc(model, &modelDesc));

  batch = modelDesc->inputTensorAttrArray->batchSize;

  if (g_batch)
  {
    batch = g_batch;
    ERROR_RET(lynModelSetDynamicBatchSizeAsync(streamSend, model, batch));
  }

  ApuInSize = modelDesc->inputDataLen;
  apuIn = new BufferPool(ApuInSize * batch, SWAP_POOL_COUNT); // n h w c
  ApuOutSize = modelDesc->outputDataLen * batch;
  apuOut = new BufferPool(ApuOutSize, SWAP_POOL_COUNT);
  if (!g_file_input.empty())
  {
    hostin = malloc(ApuInSize * batch);
    if (nullptr == hostin)
    {
      std::cerr << "malloc failed " << g_file_input.c_str();
      return -1;
    }
    std::ifstream inputFile(g_file_input, std::ios::in | std::ios::binary);
    if (false == inputFile.is_open())
    {
      std::cout << "test open failed " << g_file_input.c_str();
      return -1;
    }

    inputFile.read((char *)hostin, ApuInSize * batch);
    inputFile.close();
  }

  /**
   * start process loop
   */
  std::cout << "TK_Info: =========== execute " << g_frameCount
            << " times start..." << std::endl;
  start_time = getCurTimeStr();
  // heightNow = modelDesc->inputTensorAttrArray[0].dims[1];//nhwc
  // widthNow = modelDesc->inputTensorAttrArray[0].dims[2];//nhwc
  if (g_dynamicHeight != 0)
  {
    ERROR_RET(lynModelSetDynamicHWSizeAsync(streamSend, model, g_dynamicHeight, g_dynamicWidth));
    std::cout << std::dec << "resolution height " << g_dynamicHeight << " width "
              << g_dynamicWidth << std::endl;
  }
  frameNum = 0;
  tFrameBeg = std::chrono::steady_clock::now();
  lynRecordEvent(streamRecv, eventStart);
  while (!g_streamError)
  {
    apuOutBuffer = apuOut->Pop();
    apuInBuffer = apuIn->Pop();

    if (!g_file_input.empty())
    {
      ERROR_RET(lynMemcpyAsync(streamSend, apuInBuffer, hostin,
                               ApuInSize * batch, ClientToServer));
    }
    ERROR_RET(lynModelSendInputAsync(streamSend, model, apuInBuffer,
                                     apuOutBuffer, batch));
    ERROR_RET(lynModelRecvOutputAsync(streamRecv, model));
    /**
     * release APU output buffer
     */
    apu_cb_info *apuCallbackInfo = new apu_cb_info;
    apuCallbackInfo->InPool = apuIn;
    apuCallbackInfo->InBuffer = apuInBuffer;
    apuCallbackInfo->OutPool = apuOut;
    apuCallbackInfo->OutBuffer = apuOutBuffer;
    ERROR_RET(lynStreamAddCallback(streamRecv, apu_buffer_callback,
                                        apuCallbackInfo));
    frameNum++;

    if ((g_frameCount > 0) && frameNum >= g_frameCount)
    {
      break;
    }
  }

  ERROR_RET(lynSynchronizeStream(streamRecv));
  lynRecordEvent(streamRecv, eventStop);
  ERROR_RET(lynSynchronizeStream(streamRecv));
  ERROR_RET(lynEventElapsedTime(eventStart, eventStop, &timees));
  ERROR_RET(lynSynchronizeStream(streamSend));

  end_time = getCurTimeStr();
  tFrameSpan = tFrameEnd - tFrameBeg;
  fps = g_frameCount * batch * 1000 / timees; // tFrameSpan.count();

  std::cout << "  start_time " << start_time << std::endl;
  std::cout << "    end_time " << end_time << std::endl;
  std::cout << "TK_Info: =========== execute " << g_frameCount << " times end"
            << std::endl;
  std::cout << std::fixed << std::setprecision(6) << "TK_Info: FPS = " << fps
            << std::endl;
  std::cout << "TK_Info: chipid = " << g_deviceId << " duration "
            << timees / 1000 << "s, FPS = " << (int)fps << std::endl;

ret:
  if (apuIn)
    delete apuIn;
  if (apuOut)
    delete apuOut;

  if (!g_file_input.empty())
  {
    if (hostin)
      free(hostin);
  }

  if (model)
    lynUnloadModel(model);
  if (eventStart)
    lynDestroyEvent(eventStart);
  if (eventStop)
    lynDestroyEvent(eventStop);
  if (streamSend)
    lynDestroyStream(streamSend);
  if (streamRecv)
    lynDestroyStream(streamRecv);

  if (context)
    lynDestroyContext(context);

  return err;
}

bool argsParser(int argc, char *argv[])
{
  CLI::App app;
  // 添加选项
  app.add_option("-m", g_argModuleFile, "model path")->required();
  app.add_option("-f", g_frameCount, "frame count")->required();
  app.add_option("-i", g_file_input, "input MP4 video file path");
  app.add_option("-o", g_file_output, "output file path");
  app.add_option("-b", g_batch, "batch size")->default_val(1);
  app.add_option("-d", g_deviceId, "lynxi device id. eg.:0")->default_val(0);
  app.add_option("-w", g_dynamicWidth, "width")->default_val(0);
  app.add_option("-y", g_dynamicHeight, "height")->default_val(0);

  CLI11_PARSE(app, argc, argv);
  return true;
}

int main(int argc, char *argv[])
{
  if (!argsParser(argc, argv))
  {
    return -1;
  }

  if ((g_dynamicWidth == 0 && g_dynamicHeight != 0) || (g_dynamicWidth != 0 && g_dynamicHeight == 0))
  {
    std::cout << "error: pls input resolution height width " << std::endl;
    return -1;
  }
  
  apuinfer();

  return 0;
}
